import json
import os
from datetime import datetime
from loguru import logger

def format_prompt(
    round_num: int,
    user_command: str,
    chat_history: str,
    memory_context: str,
    file_context: str,
    tool_descriptions: str,
    current_situation: dict,
    processing_history: list,
    all_plan: list = None,
    current_step_in_plan: int = 0,
    last_info: str = None,
    tool_result:str = ""
) -> str:
    current_date = datetime.now().strftime("%Y年%m月%d日 %H:%M:%S")
    current_working_dir = os.getcwd()
    os_info = current_situation.get("os_info", "Unknown")
    tool_calls_format = "<tool_calls><tool><tool_name>...</tool_name>" \
    "<params><param name>..</param name><path>.</path> ... </params>" \
    "<expect>dir list, files or directories in the specified path</expect>" \
    "<dependencies>...</dependencies><time>30</time>" \
    "<failure_policy>retry 3 times</failure_policy></tool>" \
    "</tool_calls>\n\n  # The params example: <params><path>.</path></params>\n\n"

# <tool_name>list_directory</tool_name>
# <params><path>C:/projects/my_project</path></params>
# <expect>dir list, files or directories in the specified path</expect>
# <dependencies></dependencies><time>10</time>
# <failure_policy>retry 3 times</failure_policy>

    # Construct the last_info section if it exists
    last_info_section = ""
    if last_info:
        last_info_section = f"""## Information from Last Task (if any)
{last_info}
"""

    if all_plan and current_step_in_plan > 0:
        # Global Plan Execution & Revision Prompt
        current_objective = all_plan[current_step_in_plan - 1].get("objective", "No objective specified for this step.")
        
        prompt = f"""# Role
你是一个顶级的 AI 助手，正在按照一个预设计划执行任务，并具备评估和修正能力。

# Context
## Global Plan (`all_plan`)
<plan>
{json.dumps(all_plan, indent=2, ensure_ascii=False)}
</plan>

**You are now at step: {current_step_in_plan}**

## Available Tools
<tools>
{tool_descriptions}
</tools>

## memorys
<memorys>
{memory_context}
</memorys>

# Rules
1.  你有一个全局计划 (`<all_plan>`)。你的当前任务是执行计划中的第 {current_step_in_plan} 步。
2.  **评估上一步结果**：仔细检查 `Processing History` 中最新的 `Observation`。
    -   如果 `Observation` 包含 "Tool Error"、"Failed" 或 "Unexpected Result" 等失败标识，或者结果明显偏离计划目标，你必须修正计划。
    -   如果结果符合预期，则继续执行下一步。
3.  在 `<think>` 标签中进行思考，说明你基于评估结果的决策。
4.  根据决策，严格遵循以下格式之一输出：

    格式 A (正常执行下一步):
    <think>思考过程：上一步执行成功，结果符合预期。现在我将执行计划的第 {current_step_in_plan} 步：{current_objective}。</think>
    {tool_calls_format}

    格式 B (修正或重写计划):
    <think>思考过程：上一步执行失败/结果异常。原因在于[...说明原因...]。我需要修正计划。修正策略是[...重写整个计划/修改某几步...].</think>
    <revise_plan>
    <task>
        <objective>修正后的第一步目标</objective>
        <expect>修正后的第一步预期结果</expect>
        <dependencies>无</dependencies>
        <estimated_time_cost>30</estimated_time_cost>
        <failure_policy>重试3次</failure_policy>
    </task>
    <task>
        <objective>修正后的第二步目标</objective>
        <expect>修正后的第二步预期结果</expect>
        <dependencies>第一步</dependencies>
        <estimated_time_cost>60</estimated_time_cost>
        <failure_policy>请求用户</failure_policy>
    </task>
    </revise_plan>

    格式 C (任务完成):
    <think>思考过程：这是计划的最后一步，且已成功执行。信息已足够，我可以总结并回复了。</think>
    <reply>你的最终答案。</reply>

5.  **除非是修正计划或最终回复，否则不要输出 `<revise_plan>` 或 `<reply>` 标签。**

## Original User Input
{user_command}

## Conversation History
<conversation>
{chat_history}
</conversation>

# Your Turn
评估上一步结果，并执行或修正计划。
"""
    elif round_num == 0:
        # Initial Prompt (First Step)
        prompt = f"""# Role
你是一个顶级的 AI 助手，能够解决复杂问题。

# Context
## Current Date and Time: {current_date}
## Operating System: {os_info}
## Current Working Directory: {current_working_dir}

{last_info_section}

## Memory Context
<memorys>
{memory_context}
</memorys>

## File Context
<files>
{file_context}
</files>

## Available Tools
<tools>
{tool_descriptions}
</tools>

# Rules
1.  仔细分析用户的需求和提供的上下文信息（包括对话历史、记忆和上次任务的遗留信息）。**在回答关于用户自身的问题时，请优先从对话历史和记忆中查找信息。**
2.  首先在 `<think>` 标签中进行思考。你的思考过程应判断任务的性质：
    -   **简单问题**: 如果根据现有信息就能直接回答，请直接在 `<reply>` 标签中给出答案。
    -   **需求不明确**: 如果用户需求模糊，请在 `<clarify>` 标签中直接提问，引导用户提供更多信息。
    -   **需要额外信息**: 如果问题需要使用工具，请不要回复，而是输出 `<tool_calls>` 标签来收集信息。
3.  你的输出必须严格遵循以下格式之一：

    格式 A (直接回复):
    <think>思考过程：这是一个简单问题/需要澄清，我可以立即回复。</think>
    <reply>你的回复或引导性问题。</reply>

    格式 B (需要行动,think + tool_calls):
    <think>思考过程：这个问题需要我调用工具才能回答。</think>
    {tool_calls_format}

    格式 C (继续行动):
    <think>思考过程：根据上次的结果，我还需要...信息。</think>
    <memory><search_1>...</search_1></memory>

## Current Original User Input
{user_command}

## Conversation History
<conversation>
{chat_history}
</conversation>

# 请给出你的回复：
"""
    elif tool_result!="":
        prompt = f"""# Role
你是一个顶级的 AI 助手，当前在处理用户的需求，请根据工具的返回结果进行回复。

## Available Tools
<tools>
{tool_descriptions}
</tools>

## Memory Context
<memorys>
{memory_context}
</memorys>
 
## Conversation History
<conversation>
{chat_history}
</conversation>

## Original User Input
{user_command}

# 工具运行的结果：
<tool_result>
{tool_result}
</tool_result>

# Your Turn
根据用户的需求和工具结果进行回复：
"""
    else:
        # Iteration Prompt (Not yet generated all_plan)
        prompt = f"""# Role
你是一个顶级的 AI 助手，正在解决一个多步任务。

## Available Tools
<tools>
{tool_descriptions}
</tools>

## Memory Context
<memorys>
{memory_context}
</memorys>

# Rules
1.  回顾整个任务、你的行动历史和获得的结果。**在回答关于用户自身的问题时，请优先从对话历史和记忆中查找信息。**
2.  在 `<think>` 标签中进行思考，判断下一步行动：
    -   **任务完成**: 如果已有信息足够回答用户最初的问题，请在 `<reply>` 标签中给出最终答案。
    -   **任务复杂，需要详细规划**: 如果你预估任务还需要超过 15 个步骤才能完成，请生成一个详细的、分步的计划，并将其放入 `<all_plan>` 标签中。
    -   **继续行动**: 如果任务还未完成但不需要生成详细计划，请继续输出 `<memory>` 或 `<tool_calls>` 标签来收集更多信息。
3.  你的输出必须严格遵循以下格式之一：

    格式 A (最终回复):
    <think>思考过程：信息已足够，我可以总结并回复了。</think>
    <reply>你的最终答案。</reply>

    格式 B (生成全局计划):
    <think>思考过程：这个任务非常复杂，我需要制定一个详细的计划来确保正确执行。</think>
    <all_plan>
    1.  [第一步计划描述]
    2.  [第二步计划描述]
    ...
    </all_plan>

    格式 C (继续行动):
    <think>思考过程：根据上次的结果，我还需要...信息。</think>
    <memory><search_1>...</search_1></memory>
    
## Original User Input
{user_command}

## Conversation History
<conversation>
{chat_history}
</conversation>

# Your Turn
根据上述信息，决定并执行下一步。
"""
    logger.debug(f"Generated Prompt (Round {round_num}):\n{prompt[:500]}...") # Log first 500 chars
    return prompt
